import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras import models, layers, optimizers
import matplotlib.pyplot as plt
import numpy as np


# 定义LSTM层函数
def create_lstm(inputs, units, return_sequences):
    lstm = layers.Bidirectional(layers.LSTM(units, return_sequences=return_sequences))(inputs)
    return lstm


# 定义Dense层函数
def create_dense(inputs, units):
    dense = layers.Dense(units, kernel_regularizer=keras.regularizers.l2(0.001), activation='relu')(inputs)
    dense_dropout = layers.Dropout(0.2)(dense)
    dense_batch = layers.BatchNormalization()(dense_dropout)
    return dense, dense_dropout, dense_batch


def load():
    # Mnist数据集加载
    (x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

    # Mnist数据集简单归一化
    x_train, x_test = x_train / 255.0, x_test / 255.0
    print('Load finished!')
    return x_train, y_train, x_test, y_test


def classify(x_train, y_train, x_test, y_test, num):
    # 设置LSTM层参数
    lstm_num_layers = num[0]
    lstm_units = num[2: 2 + lstm_num_layers]
    lstm_name = list(np.zeros((lstm_num_layers,)))

    # 设置LSTM_Dense层参数
    lstm_dense_num_layers = num[1]
    lstm_dense_units = num[2 + lstm_num_layers: 2 + lstm_num_layers + lstm_dense_num_layers]
    lstm_dense_name = list(np.zeros((lstm_dense_num_layers,)))
    lstm_dense_dropout_name = list(np.zeros((lstm_dense_num_layers,)))
    lstm_dense_batch_name = list(np.zeros((lstm_dense_num_layers,)))

    inputs_lstm = layers.Input(shape=(x_train.shape[1], x_train.shape[2]))
    for i in range(lstm_num_layers):
        if i == 0:
            inputs = inputs_lstm
        else:
            inputs = lstm_name[i - 1]
        if i == lstm_num_layers - 1:
            return_sequences = False
        else:
            return_sequences = True
        lstm_name[i] = create_lstm(inputs, lstm_units[i], return_sequences)
    for i in range(lstm_dense_num_layers):
        if i == 0:
            inputs = lstm_name[lstm_num_layers - 1]
        else:
            inputs = lstm_dense_batch_name[i - 1]
        lstm_dense_name[i], lstm_dense_dropout_name[i], lstm_dense_batch_name[i] = create_dense(inputs,
                                                                                                lstm_dense_units[i])
    outputs_lstm = layers.Dense(10, activation='softmax')(lstm_dense_batch_name[lstm_dense_num_layers - 1])

    LSTM_model = keras.Model(inputs_lstm, outputs_lstm)
    LSTM_model.compile(optimizer=keras.optimizers.Adam(),
                       loss='sparse_categorical_crossentropy',
                       metrics=['accuracy'])

    history = LSTM_model.fit(x_train, y_train, batch_size=32, epochs=1, validation_split=0.1, verbose=0)
    print('LSTM finished!')

    # 验证模型：
    results = LSTM_model.evaluate(x_test, y_test, verbose=0)
    return results[1]
